# Ultralytics YOLO 🚀, AGPL-3.0 license
""" Yolo 卷积模块 Convolution modules."""

import math

import numpy as np
import torch
import torch.nn as nn

__all__ = (
    "Conv",
    "Conv2",
    "LightConv",
    "DWConv",
    "Focus",
    "GhostConv",
    "ChannelAttention",
    "SpatialAttention",
    "CBAM",
    "Concat",
#    "DWConvTranspose2d",
#    "ConvTranspose",
#    "RepConv",
)


def autopad(k, p=None, d=1):  # kernel, padding, dilation
    """计算填充尺寸。如果有定义填充尺寸，就使用定义的填充。如果没有，返回填充尺寸，使得输出的特征图形状与输入一致（不考虑步长的情况下）。
    Args:
        k (_type_): 卷积核尺寸
        p (_type_, optional): 自定义填充尺寸，如果没有，为None. Defaults to None.
        d (int, optional): 卷积核扩张率. Defaults to 1.
    Returns:
        _type_: 填充尺寸
    """
    if d > 1:
        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p


class Conv(nn.Module):
    """标准卷积模块：输入>>卷积>>规范化>>激活函数>>输出。
    """
    default_act = nn.SiLU()  # default activation

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
        """初始化卷积模块，这里是标准卷积模块： 输入>>卷积>>规范化>>激活函数>>输出。默认激活函数为SiLu。
        注意：act 也可以直接用另一个卷积层作为激活函数，但不建议这样做。建议只用nn.Module指定的激活函数。
        Args:
            c1 (_type_): 输入通道数量
            c2 (_type_): 输出通道数量
            k (int, optional): 卷积核大小. Defaults to 1.
            s (int, optional): 步长大小，用于缩小特征图尺寸：输出特征图尺寸=输入特征图尺寸/步长. Defaults to 1.
            p (_type_, optional): 填充尺寸(一般默认None，程序自动计算，使得输出特征图尺寸与输入特征图一样). Defaults to None.
            g (int, optional): 通道分组，一般为输入通道数量与输出通道数量的公约数(这样通道才能分配均匀). Defaults to 1.
            d (int, optional): 卷积核扩张率. Defaults to 1.
            act (bool, optional): 激活函数，true表示使用默认的SiLu函数。如果是nn.Module类型表示使用指定激活函数。否则无类型，直通. Defaults to True.
        """

        """补充说明：
pytorch 的卷积函数有一个参数：groups。
为了计算方便，约定 输入通道数A与输出通道数B有公约数G，令a = A/G, b = B/G .
即：输入输出通道分成G组 a通道输入b通道输出的卷积操作，用(方式1)进行计算，最后将输出拼接起来。
只需要 G*a*b=A*B/G个卷积核。有 A * B / G 次卷积运算。
优点：减少计算量，增强不同组之间的特异性。即不同的组可以安排不同的"学习任务"。  

dilation(卷积核扩张率，默认为1)。
这个参数会使得卷积核的作用范围扩大到Keff =  k+(k−1)×(d−1)。
比如3*3卷积，dilation = 2时，相当于5*5的卷积：
k00, 000, K01 , 000, K02,  
000, 000, 000, 000, 000,
k00, 000, K01 , 000, K02,  
000, 000, 000, 000, 000,
k00, 000, K01 , 000, K02,  
显然，diltation 参数会影响输出尺寸。
显然，dilation参数可以扩大卷积核感受域，又不增加计算量，不改变输入输出尺寸。
        """

        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False) # type: ignore
        self.bn = nn.BatchNorm2d(c2) 
        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()

        """计算该层模型大致的参数量。单位是：个。
        卷积核数量 = 输入通道数*输出通道数/分组数。
        所以参数量 = 卷积核数量*卷积核大小
        """
        self.cost_memory = c1*c2/g*k*k
        """计算时间复杂度。这里用大致的乘法次数衡量.
        (每次乘法运算，基本都会搭配一次加法运算)。
        激活函数当作一次乘法运算。
        首先，几乎所有的参数都要参与计算一次，才能得到一个特征图的位置的值。
        所以，花费的时间 = 输出特征图尺寸 * 参数量。
        而特征图尺寸是无法在初始化的时候得到的，前向传播时才能得到。 
        """
        # self.cost_time 

    def forward(self, x):
        """输入>>卷积>>规范化>>激活函数>>输出。
        Args:
            x (_type_): 输入特征图
        Returns:
            _type_: 输出特征图
        """
        return self.act(self.bn(self.conv(x)))

    def forward_fuse_unused(self, x):
        """fuse:保险丝，融合，使融化。这里是假设归一化层融合到卷积层里了，前向推理跳过了规范化层。
        很可惜，这里卷积self.conv的偏置项是false，无法融合。后面也未使用。
        (可以在融合的时候，重新初始化self.conv，待改进)  

        推理阶段，假设模型中的卷积层和批归一化层已经融合，
        不需要执行归一化操作操作，提高计算效率。
        输入>>卷积>>激活函数>>输出。
        Args:
            x (_type_): 输入特征图
        Returns:
            _type_: 输出特征图
        """
        return self.act(self.conv(x))


class Conv2(Conv):
    """Simplified RepConv module with Conv fusing： 
    带卷积融合的简化版重参数化卷积模块。  -- 继承自标准卷积模块。

RepConv (Re-parameterized Convolution)重参数化卷积。
核心思想：
训练时：使用多分支结构（如 3×3 卷积 + 1×1 卷积 + 恒等映射(即f(x)=x) ）增强模型表达能力； 
推理时：通过数学方法将多分支结构等价转换为单个卷积核（如一个 3×3 卷积），从而在不损失精度的前提下提升计算效率。
与批归一化类似，反向传播的时候，要适当拆解一下结构，训练效果比较好。
    """

    def __init__(self, c1, c2, k=3, s=1, p=None, g=1, d=1, act=True):
        """标准卷积模块的初始化参数。这里当作标准卷积模块使用即可，不影响结构。
        Args:
            c1 (_type_): 输入通道数量
            c2 (_type_): 输出通道数量
            k (int, optional): 卷积核大小. Defaults to 1.
            s (int, optional): 步长大小，用于缩小特征图尺寸：输出特征图尺寸=输入特征图尺寸/步长. Defaults to 1.
            p (_type_, optional): 填充尺寸(一般默认None，程序自动计算，使得输出特征图尺寸与输入特征图一样). Defaults to None.
            g (int, optional): 通道分组，一般为输入通道数量与输出通道数量的公约数(这样通道才能分配均匀). Defaults to 1.
            d (int, optional): 卷积核扩张率. Defaults to 1.
            act (bool, optional): 激活函数，true表示使用默认的SiLu函数。如果是nn.Module类型表示使用指定激活函数。否则无类型，直通. Defaults to True.
        """
        # 先初始化父类变量
        super().__init__(c1, c2, k, s, p, g=g, d=d, act=act)

        # 新定义一个卷积层。与标准卷积层的输入输出尺寸一致，只是卷积核大小为1.
        self.cv2 = nn.Conv2d(c1, c2, 1, s, autopad(1, p, d), groups=g, dilation=d, bias=False)   # type: ignore

    def forward(self, x):
        """Apply convolution, batch normalization and activation to input tensor."""
        return self.act(self.bn(self.conv(x) + self.cv2(x)))

    def forward_fuse(self, x):
        """Apply fused convolution, batch normalization and activation to input tensor."""

        """ 这个类的 fuse_convs 函数没有融合规范化层，所以还需要规范化"""
        return self.act(self.bn(self.conv(x)))

    def fuse_convs(self):
        """Fuse parallel convolutions.
        参数融合。将输入的两组并行卷积核融合
        """

        # self.conv是标准卷积层(在父类Conv中有定义)，权重数据尺寸为 [C_out, C_in, kernel_height, kernel_width]
        # C_out：输出通道数（卷积核的数量）
        # C_in：输入通道数
        # kernel_height 和 kernel_width：卷积核的高度和宽度
        w = torch.zeros_like(self.conv.weight.data)  # 这里获得一个零张量

        # w.shape 返回张量的维度元组，[2:] 取元组的第三个元素（索引 2）到最后
        i = [x // 2 for x in w.shape[2:]]  # 最后 i = [kernel_height//2, kernel_width//2]

        #  :,:, 表示选择所有输出通道和输入通道，
        #  i[0]:i[0]+1 表示选择 【i[0],i[0]+1）的位置。i[0]取，i[1]不取。（切片的规则。）
        #  i[1] : i[1] + 1 也是同理。最后得到 [C_out, C_in, 1, 1] 的张量(数据本体还在w)。 
        # 使用clone() 函数，防止后续修改 w 时影响原始的 1×1 卷积核
        w[:, :, i[0] : i[0] + 1, i[1] : i[1] + 1] = self.cv2.weight.data.clone() 
    
        self.conv.weight.data += w  
        self.__delattr__("cv2")      # 删除 1*1 卷积层，释放内存。
        self.forward = self.forward_fuse  # 修改前向传播函数，直接使用融合之后的卷积

class DWConv(Conv):
    """Depth-wise convolution. 逐通道卷积，在标准卷积模块的基础上，尽可能最大的分组：分组数为输入通道数与输出通道数的最大公约数。
    一般来这个模块的使用方法是：输入通道数等于输出通道数。
    """

    def __init__(self, c1, c2, k=1, s=1, d=1, act=True):  # ch_in, ch_out, kernel, stride, dilation, activation
        """Initialize Depth-wise convolution with given parameters. 
        初始化的参数都是标准卷积模块的参数，其中通道数为输入通道数与输出通道数的最大公约数 """
        super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)

class LightConv(nn.Module):
    """
    Light convolution with args(ch_in, ch_out, kernel).
    轻量级卷积模块。由深度分离卷积和逐点卷积(核大小为1的普通卷积)组合而成。
    逐点卷积能增加通道数量，但不能整合邻域信息。
    深度分离卷积一般不增加通道数量。但能整合邻域信息。
    因此，这两个模块组合使用，可以达到普通卷积模块的效果，同时减少计算量。
    先通过逐点卷积增加通道数量，然后通过深度分离卷积组合新的特征图。
    """

    def __init__(self, c1, c2, k=1, act=nn.ReLU()):
        """Initialize Conv layer with given arguments including activation."""
        super().__init__() 
        self.conv1 = Conv(c1, c2, 1, act=False)
        self.conv2 = DWConv(c2, c2, k, act=act) # type: ignore
        self.cost_memory = self.conv1.cost_memory+self.conv2.cost_memory   

    def forward(self, x):
        """Apply 2 convolutions to input tensor."""
        return self.conv2(self.conv1(x))



class Focus(nn.Module):
    """Focus wh information into c-space. 将宽高信息聚焦到通道空间中。
    就是将输入张量的特征图宽高减半，四等分采样得到4组通道，通道拼接起来，送到标准卷积模块。
    
优势
1. 高效特征提取
参数量减少：相比传统的步长为 2 的卷积，该操作通过无参数的采样和通道拼接实现下采样，参数量减少约 75%。
计算速度提升：避免了卷积运算中的大量乘法累加操作，适合移动设备部署。
2. 多尺度信息融合
空间信息保留：每个输出通道整合了原图像中 2×2 区域的信息，保留了局部结构（如边缘、角点）。
隐式特征增强：拼接操作使后续卷积能够同时处理不同位置的特征，增强模型对局部模式的感知能力。

    注意：PyTorch 标准：深度学习模型通常处理批量数据，输入默认格式为 [B, C, H, W]。
    """

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
        """Initializes Focus object with user defined channel, convolution, padding, group and activation values."""
        super().__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act=act)
        # self.contract = Contract(gain=2)
        self.cost_memory = self.conv.cost_memory 

    def forward(self, x):
        """
        Applies convolution to concatenated tensor and returns the output.

        Input shape is (b,c,w,h) and output shape is (b,4c,w/2,h/2).
        输入张量x ，将空间信息放到通道信息，然后通道数放大到对应的倍数。

        x[..., ::2, ::2]      
        其中 “...”  表示匹配所有未显式指定的维度。
        不管x是 [B,C,H,W] 形状还是 [C,H,W] 形状的张量，都只对最后两个维度H,W操作，[B,C]或[C]维度不变。 
        其中 “1::2” 表示从索引1开始，每隔1个元素取一个值(步长为2.)，等效于提取基数索引。

        x[..., ::2, ::2]：提取所有批次、所有通道，但只保留图像中偶数行和偶数列的像素。
        x[..., 1::2, ::2]：提取所有批次、所有通道，但只保留图像中奇数行和偶数列的像素。
        x[..., ::2, 1::2]：提取所有批次、所有通道，但只保留图像中偶数行和奇数列的像素。
        x[..., 1::2, 1::2]：提取所有批次、所有通道，但只保留图像中奇数行和奇数列的像素。

输入 x:
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
# 四等分采样
a = x[..., ::2, ::2]  # 左上角
b = x[..., 1::2, ::2]  # 左下角
c = x[..., ::2, 1::2]  # 右上角
d = x[..., 1::2, 1::2]  # 右下角


左上角采样 a:
tensor([[ 0,  2],
        [ 8, 10]])

左下角采样 b:
tensor([[ 4,  6],
        [12, 14]])

右上角采样 c:
tensor([[ 1,  3],
        [ 9, 11]])

右下角采样 d:
tensor([[ 5,  7],
        [13, 15]])
        """

#  torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1)
#  dim：指定拼接维度.比如 [B,C,H,W]形状张量，（0 表示批次维度，1 表示通道维度，2 表示高度维度，3 表示宽度维度）。
#  dim = 1 表示拼接通道维度。      

# ! >>  输入必须是 [B,C,H,W]形状张量。四等分采样：每个切片的形状为 [B, C, H/2, W/2]。拼接输出的尺寸： [B, 4C, H/2, W/2]
        return self.conv(
            torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1)
            )
        # return self.conv(self.contract(x))


class GhostConv(nn.Module):
    """Ghost Convolution https://github.com/huawei-noah/ghostnet.
    Ghost 卷积（Ghost Convolution）是华为诺亚方舟实验室在 2020 年提出的一种高效卷积结构，旨在减少参数量和计算量的同时保持模型性能。
    Ghost"（幽灵）大多数卷积生成的特征图中存在大量冗余，可以被视为 "幽灵特征"（Ghost Features）
    原论文《GhostNet: More Features from Cheap Operations》将这种方法类比为：
    "我们观察到深度卷积神经网络中的特征图包含大量冗余，这些冗余特征可以通过简单变换从 ' 真实 ' 特征中派生出来，就像幽灵从实体中产生一样。"

    Ghost 卷积的核心思想是：
        大多数卷积产生的特征图中存在大量冗余，可以通过廉价操作（如5×5深度可分离卷积）从已有特征图中派生出来。
    操作思路：
        先输出较小的通道数，生成基础特征图，
        然后用低成本操作(深度可分离卷积即逐通道卷积)生成派生特征图。
        将基础特征图与派生特征图拼接，得到最终输出。
    """

    def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
        """
        注意：输出通道数必须为2的倍数，否则会有问题。
        Initializes Ghost Convolution module with primary and cheap operations for efficient feature learning."""
        super().__init__()
        c_ = c2 // 2  # hidden channels 
        self.cv1 = Conv(c1, c_, k, s, None, g, act=act)
        self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)
        self.cost_memory = self.cv1.cost_memory+self.cv2.cost_memory

    def forward(self, x):
        """Forward propagation through a Ghost Bottleneck layer with skip connection."""
        y = self.cv1(x)
        return torch.cat((y, self.cv2(y)), 1)



class ChannelAttention(nn.Module):
    """Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet.
    
    通道注意力：
    具体步骤如下：
全局信息聚合：通过全局平均池化（Global Average Pooling）将每个通道的空间信息压缩为一个标量。
通道重要性评估：通过一个小型神经网络（通常是 1×1 卷积）学习各通道的重要性权重。
特征重标定：使用 Sigmoid 激活函数将权重归一化到 [0,1] 区间，并与原始特征相乘，增强重要通道，抑制不重要的通道。

这里是极简形式的通道注意力。
与其他通道注意力对比可以改变的方向：
1、池化还有最大池化，但平均池化对训练更有利，实验结论：平均池化在准确率和效率上取得最佳平衡。

2、这里小型神经网络（1×1卷积）只用了一层，可以用两层。不过这个模块没有残差操作，太深对训练不利。
    """

    def __init__(self, channels: int) -> None:
        """Initializes the class and sets the basic configurations and instance variables required."""
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1) # 全局平均池化。
        """nn.AdaptiveAvgPool2d(1) 是 PyTorch 中的自适应平均池化层，其作用是将输入特征图的空间维度强制压缩为 1×1。
        与普通平均池化（如 nn.AvgPool2d(kernel_size=2)）不同，自适应池化的输出尺寸由用户指定，而不是由卷积核大小决定。
        """
        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
        self.act = nn.Sigmoid()
        self.cost_memory = channels*channels

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Applies forward pass using activation on convolutions of the input, optionally using batch normalization."""
        return x * self.act(self.fc(self.pool(x)))


class SpatialAttention(nn.Module):
    """Spatial-attention module.
    空间注意力。空间注意力的核心思想是：不同空间位置对特定任务的贡献不同，应动态调整各位置的权重。

    具体步骤如下：
特征聚合：通过通道维度的平均池化和最大池化，将多通道特征压缩为两个统计特征图。
空间特征提取：通过卷积操作融合这两个统计特征图，生成空间注意力图。
特征重标定：使用 Sigmoid 激活函数将注意力图归一化到 [0,1] 区间，并与原始特征相乘，增强重要区域，抑制不重要区域。
    """

    def __init__(self, kernel_size=7):
        """Initialize Spatial-attention module with kernel size argument.
        kernel_size 只能选3 或 7. 
        实验测试卷积核等于3的时候，效果与5差不多。
        卷积核等于7时， ImageNet 分类任务中效果提高约0.6%，卷积核再大就没什么提升效果了，反而导致过拟合。
        卷积核越大，参数量与计算量就越大。 选择 3 或 7 是在模型表达能力与计算效率之间的最佳平衡点。
        """
        super().__init__()
        assert kernel_size in {3, 7}, "kernel size must be 3 or 7"  
        padding = 3 if kernel_size == 7 else 1
        self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.act = nn.Sigmoid() # Sigmoid 激活函数,将注意力图的每个值压缩到 [0,1] 区间，表示空间位置的重要性
        self.cost_memory = 2 * kernel_size*kernel_size # 空间注意力是很轻量级的操作，只有两个卷积核的参数

    def forward(self, x):
        """Apply channel and spatial attention on input for feature recalibration."""
        return x * self.act(
            self.cv1(
                torch.cat(
                    [
                        torch.mean(x, 1, keepdim=True),    
                        # [B,C,H,W] → [B,1,H,W]，在通道维度（dim=1）上计算平均值.keepdim=True 保留维度，确保输出仍为 4D 张量
                        torch.max(x, 1, keepdim=True)[0]],  
                        # [B,C,H,W] → [B,1,H,W]，torch.max 返回两个值：最大值和索引。[0]取最大值部分。keepdim=True 同样保留维度。
                    1)  # 在通道维度（dim=1）拼接平均响应和最大响应
                )
            )

class CBAM(nn.Module):
    """Convolutional Block Attention Module.
    混合注意力模块：先通道注意力，再空间注意力。
    """

    def __init__(self, c1, kernel_size=7):
        """Initialize CBAM with given input channel (c1) and kernel size."""
        super().__init__()
        self.channel_attention = ChannelAttention(c1)
        self.spatial_attention = SpatialAttention(kernel_size)

    def forward(self, x):
        """Applies the forward pass through C1 module."""
        return self.spatial_attention(self.channel_attention(x))


class Concat(nn.Module):
    """Concatenate a list of tensors along dimension."""

    def __init__(self, dimension=1):
        """Concatenates a list of tensors along a specified dimension."""
        super().__init__()
        self.d = dimension

    def forward(self, x):
        """Forward pass for the YOLOv8 mask Proto module."""
        return torch.cat(x, self.d)








